"""
Module to extract and analyze data from the Anton Paar MCR rheometer exported
the the Rheoplus software.
Created: Brent Maranzano
Last Modified: June 13. 2014
"""
import csv
import numpy as np
from scipy import delete
import h5py
from pdb import set_trace

def import_file_old(data_name, HDF_name, HDF_group):
    """
    Created: Brent Maranzano
    Last Modified: June 13, 2014
    Import the text file obtained from the tables of Rheoplus software.
    data_name : String name of txt file with rheology data.
    HDF_name: String name of HDF file to place data.
    HDF_group: String name of group within HDF file to place data.
    """
    def create_datagroup(HDF_group):
        """
        Function to create group in HDF file at the point where
        "Data Series Information" is found in the text data file.
        """
        line = reader.next()
        if line[0].find('Name:') == 0:
            group_name = line[0].split(':')[1].strip()
            group_name = '%s/%s'%(HDF_group,group_name)
            f_h5py.create_group(group_name )
            line = reader.next()
        else:
            raise ImportError
        if line[0].find('Sample:') == 0:
            f_h5py[group_name].attrs['Sample'] = line[0].split(':')[1].strip()
            line = reader.next()
        if line[0].find('Operator:') == 0:
            f_h5py[group_name].attrs['Operator'] = line[0].split(':')[1].strip()
            line = reader.next()
        if line[0].find('Remarks:') == 0:
            f_h5py[group_name].attrs['Remarks'] = line[0].split(':')[1].strip()
            line = reader.next()
        if line[0].find('Number of Intervals:') == 0:
            intervals = int(line[0].split(':')[1].strip())
            f_h5py[group_name].attrs['Intervals'] = intervals
            for i in xrange(intervals):
                create_dataset(f_h5py[group_name])
        return

    def create_dataset(hdf_group):
        """
        Function to a create dataset in HDF file at the point where
        "Interval:" is found in the text data file.
        """
        num = 0
        for line in reader:
            if line[0].find('Interval') == 0:
                interval = int(line[0].split(':')[1].strip())
                line = reader.next()
                if line[0].find('Number of Data Points') == 0:
                    num = int(line[0].split(':')[1].strip())
                else:
                    raise ImportError
            if line[0].find('Meas. Pts.') == 0:
                break
        header = line[0].split('\t')
        units = reader.next()[0].split('\t')
        data = []
        mask = []
        for n in xrange(num):
            line = reader.next()[0]
            if line.find('invalid point') == -1:
                d = []
                m = []
                for i, c in enumerate(line.split('\t')):
                    try:
                        d.append(float(c.replace(",","")))
                        m.append(1)
                    except:
                        m.append(0)
                data.append(d)
                mask.append(m)
        data = np.array(data)
        header = [header[i] for i, m in enumerate(mask[0]) if m==1]
        units = [units[i] for i, m in enumerate(mask[0]) if m==1]
        adata = np.zeros(data.shape[0], dtype={'names':header, 'formats':len(header)*['f4']})
        for i, f in enumerate(header):
            adata[f] = data[:,i]
        try:
            hdf_data = hdf_group.create_dataset("Interval %i"%interval, data=adata)
        except:
            set_trace()
        hdf_data.attrs['header'] = header
        hdf_data.attrs['units'] = units
        f_h5py.flush()
        return

    f_h5py = h5py.File(HDF_name, mode='a')
    with open(data_name, mode='r') as fObj:
        reader = csv.reader(fObj, delimiter="#")
        for line in reader:
            if line[0].find('Data Series Information') == 0:
                create_datagroup(HDF_group)

def import_file(fname, HDF_name, HDF_group):
    """
    Created: Brent Maranzano
    Last Modified: June 13, 2014
    Import the text file obtained from the tables of Rheoplus software.
    fname : String name of txt file with rheology data.
    HDF_name: String name of HDF file to place data.
    HDF_group: String name of group within HDF file to place data.
    """
    def read_file():
        fObj = open(fname, mode='r')
        # Find the positions of each Data Series
        series_loc = []
        for ln, ltxt in enumerate(iter(fObj.readline, "")):
            if ltxt.find("Data Series Information") == 0:
                series_loc.append((ln, fObj.tell()))
        series_loc.append((ln, fObj.tell())) # append the last line number in file

        all_series = []
        # Loop over each Data Series
        for i in range(len(series_loc[:-1])):
            fObj.seek(series_loc[i][1])
            series = {'Name':'', 'Sample':'', 'Operator':'', 
                    'Remarks':'', 'Number of Intervals':'', 'Interval':[]}
            count = 0
            # Populate Data Series Information
            for j in range(10):
                line = fObj.readline()
                count += 1
                for kw in series.keys():
                    if line.find(kw) == 0:
                        series[kw] = line.split(':')[1].strip()
            # Get the data for each interval in the Data Series
            while count < (series_loc[i+1][0] - series_loc[i][0]):
                line = fObj.readline()
                count += 1
                if line.find('Interval') == 0:
                    interval = {'Interval': None, 'Number Points': None,
                            'Header': None, 'Units': None, 'Data': []}
                    interval['Interval'] = line.split(':')[1].strip()
                    line = fObj.readline()
                    count += 1
                    interval['Number Points'] = int(line.split(':')[1].strip())
                    while line.find('Meas. Pts.') != 0:
                        line = fObj.readline()
                        count += 1
                    interval['Header'] = [h.strip() for h in line.split('\t')]
                    line = fObj.readline()
                    count += 1
                    interval['Units'] = [h.strip() for h in line.split('\t')]
                    for data_pt in range(interval['Number Points']):
                        line = fObj.readline()
                        count += 1
                        # remove commas from data formats
                        line = [l.strip().replace(",","") for l in line.split('\t')]
                        # if last item doesn't have "WMa" or "DSO" or ...etc., add a status
                        if (line[-1].find("WMa") == -1) and (line[-1].find("DSO") == -1):
                            line.append("WMa")
                        interval['Data'].append(line)
                    series['Interval'].append(interval)
            all_series.append(series)
        fObj.close()
        return all_series

    def clean_data(data):
        for series in data:
            for interval in series['Interval']:
                d = np.array(interval['Data'])
                mask = np.zeros(d.shape)
                for i in xrange(d.shape[0]):
                    for j in xrange(d.shape[1]):
                        try:
                            float(d[i,j])
                        except:
                            mask[i,j] = True
                d = np.ma.masked_array(d, mask=mask)
                bad_rows = []
                for i in xrange(mask.shape[0]):
                    if np.all(mask[i,1:] == 1):
                        bad_rows.append(i)
                good_cols = []
                bad_cols = []
                for j in xrange(mask.shape[1]):
                    if np.all(mask[:,j] == 1):
                        bad_cols.append(j)
                    else:
                        good_cols.append(j)
                interval['Header'] = [interval['Header'][i] for i in good_cols]
                interval['Units'] = [interval['Units'][i] for i in good_cols]
                d = delete(d, bad_rows, 0)
                d = delete(d, bad_cols, 1)
                new_data = np.zeros(d.shape[0], dtype={'names':interval['Header'], 
                    'formats':len(interval['Header'])*['f4']})
                for j in range(d.shape[1]):
                    new_data[interval['Header'][j]] = d[:,j]
                interval['Data'] = new_data
        return data

    def create_hdf(data):
        f_h5py = h5py.File(HDF_name, mode='a')
        for series in data:
            grp = '%s/%s'%(HDF_group, series['Name'])
            f_h5py.create_group(grp)
            for desc in ['Sample', 'Operator', 'Remarks', 'Number of Intervals']:
                f_h5py[grp].attrs[desc] = series[desc]
            for interval in series['Interval']:
                dst = '%s/%s'%(grp, interval['Interval'])
                f_h5py.create_dataset(dst, data=interval['Data'])
                f_h5py[dst].attrs['Units'] = interval['Units']
        f_h5py.flush()
        f_h5py.close()

    data = read_file()
    data = clean_data(data)
    create_hdf(data)
    return 

